import socket
import time
import traceback
import select
import os
from mqtt_conf import*


port = 23

# \xff\xfe\x03 \xff\xfb\x22 \xff\xfe\x01
# IAC WILL Suppress_Go_Ahead
# IAC DONT LINEMODE
# IAC DONT ECHO
# \xff\xfb\x01 \xff\xfb\x03 \xff\xfb\x22
# IAC WILL ECHO
# IAC WILL Suppress_Go_Ahead
# IAC WILL LINEMODE

IAC=255 # \xff
SB=250  # \xfa
SE=240  # \xf0
NAWS=31 # \x1f
DO=253  # \xfd
DONT=254 # \xfe
WILL=251 # \xfb
WONT=252 # \xfc
ECHO=1   # \x01
Suppress_Go_Ahead=3 # \x03
LINEMODE=34 # \x22



def filter_telnet(b: bytes):
    n = len(b)
    res = bytearray(n)
    j = 0
    i = 0

    while i < n:
        if b[i] != IAC:
            res[j] = b[i]
            j += 1
            i += 1
            continue
        i += 1
        if b[i]==SB:
            i += 1
            if b[i]==NAWS: # NAWS 0 col col/255 row row/255 IAC SE
                if i+7 >= n:
                    i+=1
                cols = b[i] * 256 + b[i+1]
                rows = b[i+2] * 256 + b[i+3]
                # mqtt_out(f'\nstty rows {rows} cols {cols}')
                client_publish(f'ext', f'size {rows} {cols}')
                # client_sock.send(bytes([IAC, WILL, SB_NAWS, IAC, SE]))
                i += 6
            continue
        if b[i]==WILL and b[i+1]==NAWS:
            # WILL NAWS IAC SB NAWS col col/256 row row/256 IAC SE
            i += 2
            continue
        if (b[i] == DONT or b[i]==DO) and b[i+1] == ECHO:
            if b[i]==DO:
                client_publish(f'ext', f'echo')
                client_sock.send(bytes([IAC, WILL, ECHO]))
            else:
                client_publish(f'ext', f'-echo')
                client_sock.send(bytes([IAC, WONT, ECHO]))
            i += 2
            continue
        if 251 <= b[i] <= 254:
            i += 2
            continue
        if b[i] == 255:
            res[j] = b[i]
            j += 1
        i += 1
    return res[:j]


def redirect_on_message(msg, topic=None):
    global client_sock
    if topic.endswith('msg'):
        if client_sock:
            client_sock.send(msg)
    elif topic.endswith('pong'):
        now_ms = time.time() * 1000
        delay_ms = f'\x1b[0;32m TTL: {now_ms - float(msg)} ms \x1b[0m\r\n'
        client_sock.send(delay_ms.encode())


def go_listen_socket(*args):
    c, addr = args
    while 1:
        has_data, _, _ = select.select([c], [], [])
        if not has_data:
            return
        b = c.recv(4096)
        if not b:
            return
        raw = filter_telnet(b)
        if not raw:
            continue
        raw = bytes(raw)
        mqtt_out_raw(raw)


if __name__ == '__main__':
    import sys
    import mqtt_conf
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    only_once = False
    if len(sys.argv) > 1 and sys.argv[1] == '-h':
        print('jc-mtc.py [-o] [port]: start a telnet server')
        print('  -o: run only once')
        print('  port: default 23')
        exit(0)

    if '-o' in sys.argv:
        only_once = True
        i = sys.argv.index('-o')
        sys.argv = sys.argv[:i] + sys.argv[i+1:]

    if len(sys.argv) > 1:
        port = int(sys.argv[1])

    ser_soc = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    ser_soc.bind(('localhost', port))
    ser_soc.listen(1)
    while 1:
        try:
            client_sock, addr = ser_soc.accept()
        except:
            ser_soc.close()
            print('accept failed')
            exit(-1)
        read_config('mqtt.ini')
        set_on_message(redirect_on_message)
        start_loop()
        now_ms = time.time() * 1000
        mqtt_out(f'pty -r -t xterm-256color')
        client_publish('ping', f'{now_ms}')
        try:
            iac_do_naws = bytes([IAC, DO, NAWS])
            client_sock.send(iac_do_naws)
            client_sock.send(f'\x1b[0;32m target:{mqtt_conf.target}\x1b[0m\r\n'.encode())
            go_listen_socket(client_sock, addr)
        except:
            traceback.print_exc()
        finally:
            client_sock.close()
            client_publish('ext', 'exit')
            stop_loop()
        if only_once:
            break
    ser_soc.close()
